/***************************************************************************
 *   Copyright (C) 2010 by Oleg Goncharov  *
 *   $EMAIL$                           *                          
 *                                                                         *
 *   This file is part of ChessVision.                                     *
 *                                                                         *
 *   ChessVision is free software; you can redistribute it and/or modify   *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/
#include "pointpyrmatch.h"

#include <opencv/cv.h>
#include <opencv/highgui.h>
#include <cmath>
#include <iostream>

using namespace std;

CPointPyrMatch::CPointPyrMatch(const cv::Mat& sample, const cv::Mat& test_img, double threshold) {
	d_maxlevel = floor(log2(MAX(sample.cols, sample.rows))) - 2;
	if (d_maxlevel < 0) d_maxlevel = 0;
	
	d_sample = new cv::Mat[d_maxlevel + 1];
	d_sample[0] = sample.clone();
	BuildPyr(d_sample);
	
	d_threshold = new double[d_maxlevel + 1];
	d_img = new cv::Mat[d_maxlevel + 1];
	
	WCShiftAnalysis(test_img);
	for(int i = 0; i <= d_maxlevel; i++) d_threshold[i] *= threshold; 
}

/**
 * @fn void WCShiftAnalysis(const cv::Mat& test_img)
 * Perform worst-case analysis pyramid thresholds tuning. It writes results to d_threshold and may decrease d_maxlevel.
 */
void CPointPyrMatch::WCShiftAnalysis(const cv::Mat& test_img) {
	cv::Mat corr;
	cv::Size shifted_size(test_img.cols -(1<<d_maxlevel) + 1, test_img.rows - (1<<d_maxlevel) + 1);
	double min, max;
	
	for(int level = 0; level <= d_maxlevel; level++) d_threshold[level] = 1.0;
	
	for(int dx = (1<<d_maxlevel) - 1; dx >=0 ; dx--) 
		for(int dy = (1<<d_maxlevel) - 1; dy >=0 ; dy--) {
			d_img[0] = cv::Mat(test_img, cv::Rect(dx, dy, shifted_size.width, shifted_size.height));
			BuildPyr(d_img);
			
			for(int level = 1; level <= d_maxlevel; level++) {
				if (dx < 1<<level && dy < 1<<level) {
					cv::matchTemplate(d_img[level], d_sample[level], corr, CV_TM_CCORR_NORMED);
					cv::minMaxLoc(corr, &min, &max);
					if (max < 0.2) d_maxlevel = level - 1;
					if (max < d_threshold[level]) d_threshold[level] = max;
				}
			}
		}
	//TODO Maybe destruct d_img[>d_maxlevel)?
}

CPointPyrMatch::~CPointPyrMatch() {
	delete [] d_sample;
	delete [] d_img;
	delete [] d_threshold;
}

bool CPointPyrMatch::DetectPoint(const IplImage * img, cv::Point2f& point) {
	d_img[0] = cv::Mat(img);
	BuildPyr(d_img);
	
	//cv::Mat corr(d_img[d_maxlevel].size()-d_sample[d_maxlevel].size()+cv::Size(1,1), CV_32FC1);
	cv::Mat corr;
	cv::Point min_loc, max_loc, best_max_loc;
	cv::Size sample_size = d_sample[d_maxlevel].size();
	double min, max, best_max = 0.0;
	
	cv::matchTemplate(d_img[d_maxlevel], d_sample[d_maxlevel], corr, CV_TM_CCORR_NORMED);
	
	for(int attempt = 0;  attempt < 50; attempt++) {
		do {
			cv::minMaxLoc(corr, &min, &max, &min_loc, &max_loc);
			if (max < d_threshold[d_maxlevel]) {
				if (best_max < d_threshold[0]) return false;
				
				point.x = best_max_loc.x + d_sample[0].cols/2;
				point.y = best_max_loc.y + d_sample[0].rows/2;	
				return true;
			}
				
			std::cout << "maxlevel max " << max << max_loc << std::endl;
			
			cv::Rect rect(max_loc, sample_size);
			
			max = RefineMatch(max_loc);
			
			rect -= cv::Point(sample_size.width/2, sample_size.height/2);
			rect &= cv::Rect(0, 0, corr.cols, corr.rows);
			std::cout << "[" << rect.x << " " << rect.y << " " << rect.width << " " << rect.height << std::endl;
			corr(rect) = 0.0;
		} while (max < d_threshold[0]);
		std::cout << "0 max " << max << max_loc << std::endl;
		if (best_max < max) {
			best_max = max;
			best_max_loc = max_loc;
		}
	}
	
	std::cout << "best max " << best_max << std::endl;
	
	point.x = best_max_loc.x + d_sample[0].cols/2;
	point.y = best_max_loc.y + d_sample[0].rows/2;	
	return true;
}

void CPointPyrMatch::BuildPyr(cv::Mat * pyr) {
	for(int i = 1; i <= d_maxlevel; i++) cv::pyrDown(pyr[i-1], pyr[i]);
}

double CPointPyrMatch::RefineMatch(cv::Point& point) {
	int d_border = 2;
	cv::Size prox_size(2*d_border + 1, 2*d_border + 1);
	cv::Size delta_size;
	cv::Rect subimg_rect;
	cv::Point min_loc, max_loc;
	cv::Mat corr;
	double min, max;
	
	for(int level = d_maxlevel - 1; level >= 0; level--) {
		point.x *= 2; point.y *= 2;
		
		delta_size = d_sample[level].size() - cv::Size(1,1);
		
		subimg_rect = cv::Rect(point, prox_size + delta_size);
		subimg_rect -= cv::Point(d_border, d_border);
		subimg_rect &= cv::Rect(cv::Point(0,0), d_img[level].size());
		
		cv::matchTemplate(d_img[level](subimg_rect), d_sample[level], corr, 
		CV_TM_CCORR_NORMED);
		
		cv::minMaxLoc(corr, &min, &max, &min_loc, &max_loc);
		if (max < d_threshold[level]) return 0.0;
		
		point = max_loc + subimg_rect.tl();
	}
	return max;
}
